# coding=utf-8
"""
聚类所有图像的所有目标宽高，分为了9类
anchor box的作用就是提供bbox的wh参考
"""
import matplotlib.pyplot as plt
import numpy as np
import os, cv2
import xml.etree.ElementTree as ET
import seaborn as sns

LABELS = ['shrinkage','blow','slag']

train_image_folder = "E:\\CAST_PNG\\CAST_Train\\CAST_selected\\"
train_annot_folder = "E:\\CAST_PNG\\CAST_Train\\Annotations\\"

"""

"""
def parse_annotation(ann_dir, img_dir, labels=[]):
    ''''
    输出
    训练图像的每个元素（dict），其内容为每一个图像的注释信息
    seen_train_labels是一个dict，其内容为
    键：目标的类
    值：在图像中发现的该目标的数目
    '''
    all_imgs = []
    seen_labels = {}
    #os.listdir(ann_dir)，方法用于返回指定的文件夹(ann_dir)包含的文件或文件夹的名字的列表
    #以字符串形式返回
    for ann in sorted(os.listdir(ann_dir)):
        #如果不是.xml文件就跳过
        if "xml" not in ann:
            continue
        img = {'object': []}
        #解析XML
        tree = ET.parse(ann_dir + ann)
        for elem in tree.iter():
            #图像的路径
            if 'filename' in elem.tag:
                #path_to_image为annotation对应的图像路径
                path_to_image = img_dir + elem.text
                img['filename'] = path_to_image
                ## 检查图像路径是否存在
                if not os.path.exists(path_to_image):
                    assert False, "file does not exist!\n{}".format(path_to_image)
            #图像的宽度
            if 'width' in elem.tag:
                img['width'] = int(elem.text)
            #图像的高度
            if 'height' in elem.tag:
                img['height'] = int(elem.text)
            #图像的标记信息，存储在obj中
            if 'object' in elem.tag or 'part' in elem.tag:
                obj = {}
                #便利object标签的子标签
                for attr in list(elem):
                    # 目标名称
                    if 'name' in attr.tag:
                        obj['name'] = attr.text

                        if len(labels) > 0 and obj['name'] not in labels:
                            break
                        else:
                            img['object'] += [obj]

                        #标签seen_labels
                        #如果在seen_labels中有此目标，则seen_labels中该目标个数加一
                        #否则初始化该目标个数为1
                        if obj['name'] in seen_labels:
                            seen_labels[obj['name']] += 1
                        else:
                            seen_labels[obj['name']] = 1
                    #图像中标记的边框信息
                    if 'bndbox' in attr.tag:
                        for dim in list(attr):
                            if 'xmin' in dim.tag:
                                #round，四舍五入
                                obj['xmin'] = int(round(float(dim.text)))
                            if 'ymin' in dim.tag:
                                obj['ymin'] = int(round(float(dim.text)))
                            if 'xmax' in dim.tag:
                                obj['xmax'] = int(round(float(dim.text)))
                            if 'ymax' in dim.tag:
                                obj['ymax'] = int(round(float(dim.text)))

        if len(img['object']) > 0:
            all_imgs += [img]
    return all_imgs, seen_labels


## Parse annotations
train_image, seen_train_labels = parse_annotation(train_annot_folder, train_image_folder, labels=LABELS)

#测试

# print("N train = {}".format(len(train_image)))
# print(train_image[2])
#
# y_pos = np.arange(len(seen_train_labels))
# fig = plt.figure(figsize=(13,10))
# ax = fig.add_subplot(1,1,1)
# ax.barh(y_pos,list(seen_train_labels.values()))
# ax.set_yticks(y_pos)
# ax.set_yticklabels(list(seen_train_labels.keys()))
# ax.set_title("The total number of objects = {} in {} images".format(
#     np.sum(list(seen_train_labels.values())),len(train_image)
# ))
# plt.show()

"""
归一化标记框大小，标记框的宽高除以图像的宽高
"""
def normalize():
    wh = []
    for anno in train_image:
        #图像的宽度
        aw = float(anno['width'])
        #图像的高度
        ah = float(anno['height'])

        for obj in anno["object"]:
            #归一化标记框的宽高，标记框宽/图像宽，标记框高/图像高
            w = (obj["xmax"] - obj["xmin"])/aw
            h = (obj["ymax"] - obj["ymin"])/ah
            temp = [w,h]
            wh.append(temp)
    wh = np.array(wh)
    return wh

"""
计算IOU
计算的是选框宽高与其聚类中心宽高的IOU值
"""
def iou(box, clusters):
    '''
    :param box:      np.array of shape (2,) containing w and h
    :param clusters: np.array of shape (N cluster, 2)
    '''
    x = np.minimum(clusters[:, 0], box[0])
    y = np.minimum(clusters[:, 1], box[1])

    intersection = x * y
    box_area = box[0] * box[1]
    cluster_area = clusters[:, 0] * clusters[:, 1]

    iou_ = intersection / (box_area + cluster_area - intersection)

    return iou_

#K-Means算法

def kmeans(boxes, k, dist=np.median, seed=1):
    #rows:标注的边框总数
    rows = boxes.shape[0]
    #记录每一个边框与聚类中心的距离。所以是形状是rows*k
    distances = np.empty((rows, k))
    #每一个边框的最新聚类中心
    last_clusters = np.zeros((rows,))
    #随机种子
    np.random.seed(seed)
    # initialize the cluster centers to be k items
    #从rows中，以概率p（该处未指定p，所以是一致分布）随机选择k个数
    #也就是随机初始化聚类中心点
    clusters = boxes[np.random.choice(rows, k, replace=False)]

    while True:
        # Step 1: allocate each item to the closest cluster centers
        #第一步，将每一项分配到最近的聚类中心
        for icluster in range(k):  # I made change to lars76's code here to make the code faster
            distances[:, icluster] = 1 - iou(clusters[icluster], boxes)

        #每个边框所属的聚类中心，distances中最小值的下标
        nearest_clusters = np.argmin(distances, axis=1)
        if (last_clusters == nearest_clusters).all():
            break

        # Step 2: calculate the cluster centers as mean (or median) of all the cases in the clusters.
        #重新计算聚类中心
        for i in range(k):
            # 每一个目标有其所属的聚类中心,nearest_clusters记载了所有选框的聚类中心（0~9）
            #i=0时，重新计算属于第一个聚类中心所有的目标的中位数
            clusters[i] = dist(boxes[nearest_clusters ==i], axis=0)

        last_clusters = nearest_clusters

    #返回值
    #cluster 聚类中心坐标
    #nearest_clusters 每个边框所属的聚类中心(1~9)
    #distance 每个边框与不同聚类中心的距离,40138*8
    return clusters, nearest_clusters, distances

# 主函数

def mian():
    kmax = 9
    dist = np.mean
    results = {}
    wh =normalize()

    #归一化后的缺陷类别分布情况
    # plt.figure(figsize=(10, 10))
    # plt.scatter(wh[:, 0], wh[:, 1], alpha=0.3)
    # plt.title("Clusters", fontsize=20)
    # plt.xlabel("normalized width", fontsize=20)
    # plt.ylabel("normalized height", fontsize=20)
    # plt.show()

    for k in range(2,kmax):
        clusters, nearest_clusters, distances = kmeans(wh,k,seed=2,dist=dist)
        WithinClusterMeanDist = np.mean(distances[np.arange(distances.shape[0]),nearest_clusters])
        result = {"clusters":             clusters,
                  "nearest_clusters":     nearest_clusters,
                  "distances":            distances,
                  "WithinClusterMeanDist": WithinClusterMeanDist}
        print("{:2.0f} clusters: mean IoU = {:5.4f}".format(k,1-result["WithinClusterMeanDist"]))
        results[k] = result

    #聚类中心乘以图像大小并保留3位小数
    clusters_round=np.around(clusters*1024,decimals=3)
    #将聚类中心展成一维数组
    clusters_list=np.array(clusters_round,dtype=np.float).ravel()

    #将结果打印在cast_anchors.txt中
    cast_achors_path="../data/anchors/cast_anchors.txt"
    if os.path.exists(cast_achors_path):os.remove(cast_achors_path)
    with open(cast_achors_path,'w') as f :
        for anchor in clusters_list:
            f.write(str(anchor)+",")
    #测试
    print(len(clusters_list))

if __name__ == '__main__':
    mian()





